package com.boulderai.mcp.server;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import org.java_websocket.WebSocket;
import org.java_websocket.handshake.ClientHandshake;
import org.java_websocket.server.WebSocketServer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.InetSocketAddress;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 简单的MCP服务器实现
 * 提供基本的工具和资源供MCP客户端使用
 */
public class SimpleMcpServer extends WebSocketServer {
    
    private static final Logger logger = LoggerFactory.getLogger(SimpleMcpServer.class);
    private final ObjectMapper objectMapper = new ObjectMapper();
    private final AtomicLong requestIdCounter = new AtomicLong(1);
    private final Map<WebSocket, Boolean> initializedClients = new HashMap<>();
    
    public SimpleMcpServer(int port) {
        super(new InetSocketAddress(port));
    }
    
    @Override
    public void onOpen(WebSocket conn, ClientHandshake handshake) {
        logger.info("New MCP client connected: {}", conn.getRemoteSocketAddress());
        initializedClients.put(conn, false);
    }
    
    @Override
    public void onClose(WebSocket conn, int code, String reason, boolean remote) {
        logger.info("MCP client disconnected: {} (code: {}, reason: {})", 
                   conn.getRemoteSocketAddress(), code, reason);
        initializedClients.remove(conn);
    }
    
    @Override
    public void onMessage(WebSocket conn, String message) {
        try {
            JsonNode messageNode = objectMapper.readTree(message);
            handleMessage(conn, messageNode);
        } catch (Exception e) {
            logger.error("Error processing message: {}", e.getMessage());
            sendError(conn, "Invalid message format", null);
        }
    }
    
    @Override
    public void onError(WebSocket conn, Exception ex) {
        logger.error("WebSocket error for client {}: {}", 
                    conn != null ? conn.getRemoteSocketAddress() : "unknown", ex.getMessage());
    }
    
    @Override
    public void onStart() {
        logger.info("MCP Server started successfully");
    }
    
    private void handleMessage(WebSocket conn, JsonNode message) {
        String method = message.has("method") ? message.get("method").asText() : null;
        JsonNode id = message.get("id");
        JsonNode params = message.get("params");
        
        if (method == null) {
            sendError(conn, "Missing method field", id);
            return;
        }
        
        switch (method) {
            case "initialize":
                handleInitialize(conn, id, params);
                break;
            case "tools/list":
                handleListTools(conn, id);
                break;
            case "tools/call":
                handleCallTool(conn, id, params);
                break;
            case "resources/list":
                handleListResources(conn, id);
                break;
            case "resources/read":
                handleReadResource(conn, id, params);
                break;
            default:
                sendError(conn, "Unknown method: " + method, id);
        }
    }
    
    private void handleInitialize(WebSocket conn, JsonNode id, JsonNode params) {
        try {
            // 创建服务器信息响应
            ObjectNode result = objectMapper.createObjectNode();
            result.put("protocolVersion", "2024-11-05");
            
            ObjectNode serverInfo = objectMapper.createObjectNode();
            serverInfo.put("name", "Simple MCP Server");
            serverInfo.put("version", "1.0.0");
            result.set("serverInfo", serverInfo);
            
            ObjectNode capabilities = objectMapper.createObjectNode();
            capabilities.put("tools", true);
            capabilities.put("resources", true);
            result.set("capabilities", capabilities);
            
            sendResponse(conn, id, result);
            initializedClients.put(conn, true);
            logger.info("Client initialized: {}", conn.getRemoteSocketAddress());
        } catch (Exception e) {
            sendError(conn, "Initialization failed: " + e.getMessage(), id);
        }
    }
    
    private void handleListTools(WebSocket conn, JsonNode id) {
        if (!isInitialized(conn)) {
            sendError(conn, "Client not initialized", id);
            return;
        }
        
        try {
            ArrayNode tools = objectMapper.createArrayNode();
            
            // 添加示例工具
            ObjectNode echoTool = objectMapper.createObjectNode();
            echoTool.put("name", "echo");
            echoTool.put("description", "回显输入的消息");
            ObjectNode echoInputSchema = objectMapper.createObjectNode();
            echoInputSchema.put("type", "object");
            ObjectNode echoProperties = objectMapper.createObjectNode();
            ObjectNode messageProperty = objectMapper.createObjectNode();
            messageProperty.put("type", "string");
            messageProperty.put("description", "要回显的消息");
            echoProperties.set("message", messageProperty);
            echoInputSchema.set("properties", echoProperties);
            ArrayNode echoRequired = objectMapper.createArrayNode();
            echoRequired.add("message");
            echoInputSchema.set("required", echoRequired);
            echoTool.set("inputSchema", echoInputSchema);
            tools.add(echoTool);
            
            // 添加计算工具
            ObjectNode addTool = objectMapper.createObjectNode();
            addTool.put("name", "add");
            addTool.put("description", "计算两个数字的和");
            ObjectNode addInputSchema = objectMapper.createObjectNode();
            addInputSchema.put("type", "object");
            ObjectNode addProperties = objectMapper.createObjectNode();
            ObjectNode aProperty = objectMapper.createObjectNode();
            aProperty.put("type", "number");
            aProperty.put("description", "第一个数字");
            ObjectNode bProperty = objectMapper.createObjectNode();
            bProperty.put("type", "number");
            bProperty.put("description", "第二个数字");
            addProperties.set("a", aProperty);
            addProperties.set("b", bProperty);
            addInputSchema.set("properties", addProperties);
            ArrayNode addRequired = objectMapper.createArrayNode();
            addRequired.add("a");
            addRequired.add("b");
            addInputSchema.set("required", addRequired);
            addTool.set("inputSchema", addInputSchema);
            tools.add(addTool);
            
            ObjectNode result = objectMapper.createObjectNode();
            result.set("tools", tools);
            
            sendResponse(conn, id, result);
        } catch (Exception e) {
            sendError(conn, "Failed to list tools: " + e.getMessage(), id);
        }
    }
    
    private void handleCallTool(WebSocket conn, JsonNode id, JsonNode params) {
        if (!isInitialized(conn)) {
            sendError(conn, "Client not initialized", id);
            return;
        }
        
        try {
            String toolName = params.get("name").asText();
            JsonNode arguments = params.get("arguments");
            
            ObjectNode result = objectMapper.createObjectNode();
            ArrayNode content = objectMapper.createArrayNode();
            
            switch (toolName) {
                case "echo":
                    String message = arguments.get("message").asText();
                    ObjectNode echoContent = objectMapper.createObjectNode();
                    echoContent.put("type", "text");
                    echoContent.put("text", "Echo: " + message);
                    content.add(echoContent);
                    break;
                    
                case "add":
                    double a = arguments.get("a").asDouble();
                    double b = arguments.get("b").asDouble();
                    double sum = a + b;
                    ObjectNode addContent = objectMapper.createObjectNode();
                    addContent.put("type", "text");
                    addContent.put("text", String.format("%.2f + %.2f = %.2f", a, b, sum));
                    content.add(addContent);
                    break;
                    
                default:
                    sendError(conn, "Unknown tool: " + toolName, id);
                    return;
            }
            
            result.set("content", content);
            result.put("isError", false);
            
            sendResponse(conn, id, result);
        } catch (Exception e) {
            sendError(conn, "Tool call failed: " + e.getMessage(), id);
        }
    }
    
    private void handleListResources(WebSocket conn, JsonNode id) {
        if (!isInitialized(conn)) {
            sendError(conn, "Client not initialized", id);
            return;
        }
        
        try {
            ArrayNode resources = objectMapper.createArrayNode();
            
            // 添加示例资源
            ObjectNode textResource = objectMapper.createObjectNode();
            textResource.put("uri", "text://example.txt");
            textResource.put("name", "示例文本");
            textResource.put("description", "一个示例文本资源");
            textResource.put("mimeType", "text/plain");
            resources.add(textResource);
            
            ObjectNode jsonResource = objectMapper.createObjectNode();
            jsonResource.put("uri", "json://config.json");
            jsonResource.put("name", "配置文件");
            jsonResource.put("description", "JSON配置文件示例");
            jsonResource.put("mimeType", "application/json");
            resources.add(jsonResource);
            
            ObjectNode result = objectMapper.createObjectNode();
            result.set("resources", resources);
            
            sendResponse(conn, id, result);
        } catch (Exception e) {
            sendError(conn, "Failed to list resources: " + e.getMessage(), id);
        }
    }
    
    private void handleReadResource(WebSocket conn, JsonNode id, JsonNode params) {
        if (!isInitialized(conn)) {
            sendError(conn, "Client not initialized", id);
            return;
        }
        
        try {
            String uri = params.get("uri").asText();
            
            ArrayNode contents = objectMapper.createArrayNode();
            
            if (uri.equals("text://example.txt")) {
                ObjectNode textContent = objectMapper.createObjectNode();
                textContent.put("uri", uri);
                textContent.put("mimeType", "text/plain");
                textContent.put("text", "这是一个示例文本文件的内容。\n包含多行文本。");
                contents.add(textContent);
            } else if (uri.equals("json://config.json")) {
                ObjectNode jsonContent = objectMapper.createObjectNode();
                jsonContent.put("uri", uri);
                jsonContent.put("mimeType", "application/json");
                ObjectNode configData = objectMapper.createObjectNode();
                configData.put("name", "示例配置");
                configData.put("version", "1.0.0");
                configData.put("enabled", true);
                jsonContent.put("text", objectMapper.writeValueAsString(configData));
                contents.add(jsonContent);
            } else {
                sendError(conn, "Resource not found: " + uri, id);
                return;
            }
            
            ObjectNode result = objectMapper.createObjectNode();
            result.set("contents", contents);
            
            sendResponse(conn, id, result);
        } catch (Exception e) {
            sendError(conn, "Failed to read resource: " + e.getMessage(), id);
        }
    }
    
    private boolean isInitialized(WebSocket conn) {
        return initializedClients.getOrDefault(conn, false);
    }
    
    private void sendResponse(WebSocket conn, JsonNode id, ObjectNode result) {
        try {
            ObjectNode response = objectMapper.createObjectNode();
            response.put("jsonrpc", "2.0");
            if (id != null) {
                response.set("id", id);
            }
            response.set("result", result);
            
            String responseJson = objectMapper.writeValueAsString(response);
            conn.send(responseJson);
            logger.debug("Sent response: {}", responseJson);
        } catch (Exception e) {
            logger.error("Failed to send response: {}", e.getMessage());
        }
    }
    
    private void sendError(WebSocket conn, String message, JsonNode id) {
        try {
            ObjectNode response = objectMapper.createObjectNode();
            response.put("jsonrpc", "2.0");
            if (id != null) {
                response.set("id", id);
            }
            
            ObjectNode error = objectMapper.createObjectNode();
            error.put("code", -1);
            error.put("message", message);
            response.set("error", error);
            
            String responseJson = objectMapper.writeValueAsString(response);
            conn.send(responseJson);
            logger.debug("Sent error: {}", responseJson);
        } catch (Exception e) {
            logger.error("Failed to send error: {}", e.getMessage());
        }
    }
    
    public static void main(String[] args) {
        int port = args.length > 0 ? Integer.parseInt(args[0]) : 9000;
        
        SimpleMcpServer server = new SimpleMcpServer(port);
        server.start();
        
        logger.info("Simple MCP Server started on port {}", port);
        logger.info("WebSocket URL: ws://localhost:{}", port);
        logger.info("Press Ctrl+C to stop the server");
        
        // 添加关闭钩子
        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            logger.info("Shutting down MCP Server...");
            try {
                server.stop();
            } catch (Exception e) {
                logger.error("Error stopping server: {}", e.getMessage());
            }
        }));
    }
}